"""
上传视频，执行目标检测，输出结果视频和目标图片
"""
import os
import subprocess
import gradio as gr
import time
from moviepy.editor import VideoFileClip

OUTPUT_DIR = os.path.join(os.getcwd(), "runs/detect")
base_conf, base_iou = 0.8, 0.2

def detect(file_path, conf, iou):
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    exp = f"exp_{timestamp}"
    output_folder = os.path.join(OUTPUT_DIR, exp)

    command = ["python", "detect.py", "--weights", "yolov5s.pt", "--source", file_path,
        "--conf-thres", str(conf), #超过这个阈值的目标会显示
        "--iou-thres", str(iou), #值越低框越少
        "--project", OUTPUT_DIR, "--name", exp,
    ]

    try:
        print("start detect using yolo-v5")
        process = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )

        stdout, stderr = process.communicate()
        print(stdout)
        print(stderr)

        output_path = os.path.join(output_folder, os.path.basename(file_path))
        print(f"yolo detect video path:{output_path}")
        final_path = os.path.join(output_folder, 'fixed_' + os.path.basename(file_path))
        print(f"moviepy convert video path:{final_path}")

        try:
            clip = VideoFileClip(output_path)
            clip.write_videofile(final_path)
        except Exception as e:
            print(f"exec moviepy convert failed: {e}")

        file_list = os.listdir(os.path.join(output_folder, 'cropped_objects'))
        print(f"exec opencv extract counts:{len(file_list)}")

        images = []
        count = 0
        for file in file_list:
            images.append(os.path.join(output_folder, 'cropped_objects', file))
            count += 1
            if count > 10:
                break

        return final_path, images

    except subprocess.CalledProcessError as e:
        print(f"detect failed:{e}")
        return None
    except Exception as e:
        print(f"detect exception:{e}")
        return None


css = """
footer {display: none !important;}
"""

demo = gr.Interface(
    inputs=[gr.Video(),gr.Slider(minimum=0, maximum=1, value=base_conf,label='target confidence show threshold'),gr.Slider(minimum=0, maximum=1, value=base_iou,label='NMS IOU show threshold')],
    outputs=[gr.Video(),gr.Gallery()],
    fn=detect,
    title="where is my paper?",
)

# demo.launch(css=css)
demo.launch(css=css, server_name="0.0.0.0")
